# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import ujson as json
import functools

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.okex.book_builder import BookBuilder
from coin.exchange.okex.kr_rest.product import OkexProduct


def _parse_side(side):
  assert side == 'bid' or side == 'ask'
  side = TradeSide.TRADE_BUY_SIDE if side == 'bid' else TradeSide.TRADE_SELL_SIDE
  return side


@functools.lru_cache(maxsize=128)
def get_product_from_channel_name(channel):
  symbol = channel[12:-6]
  return OkexProduct.FromStrNativeProduct(symbol)


class OkexHandler(WebsocketHandler):
  def on_first_book(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())

    channel = message['channel']
    product = get_product_from_channel_name(channel)
    native_symbol = product.native_symbol
    if not self.is_subscribed('book', native_symbol):
      return
    data = message['data']

    book = BookBuilder(native_symbol, book_builder_observer=self._book_builder_observer)
    bids = data.get('bids', [])
    asks = data.get('asks', [])
    book.update_fullbook(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self._books[native_symbol] = book

    self.ready = True
    self.publish_event('book', book.symbol, book)
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)

  def on_diff(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())

    channel = message['channel']
    product = get_product_from_channel_name(channel)
    native_symbol = product.native_symbol
    if native_symbol not in self._books:
      return

    data = message['data']

    book = self._books[native_symbol]
    bids = data.get('bids', [])
    asks = data.get('asks', [])
    book.update_diff(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    message = record.data
    if not isinstance(message, dict) and not isinstance(message, list):
      message = json.loads(message.tobytes())

    channel = message['channel']
    assert channel.startswith('ok_sub_spot_')
    assert channel.endswith('deals')
    product = get_product_from_channel_name(channel)
    native_symbol = product.native_symbol

    if not self.is_subscribed('trade', native_symbol):
      return
    data = message['data']
    if not isinstance(data, list):
      # {'binary': 0, 'channel': 'ok_sub_spot_eos_bch_deals', 'data': {'result': False, 'error_msg': 'The subscription event type is invalid.', 'error_code': 20103}}
      # This kind of response happens if the pair is delisted.
      return

    for fill_msg in data:
      price = float(fill_msg[1])
      qty = float(fill_msg[2])
      side = _parse_side(fill_msg[4])
      trade_event = create_market_trade_event(record.timestamp,
                                              native_symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)
      self.publish_event('trade', native_symbol, trade_event)
